#include <vector>

#include "ai_service_impl.h"
#include "ai_service_log.h"
#include "ai_service_errors.h"
#include "YOLOPostProcess.h"
#include <iostream>
#include <fstream>
#include <sys/time.h>

using namespace OHOS::HiviewDFX;
using namespace cv;

namespace OHOS {
namespace AIS {
namespace {
    constexpr HiLogLabel LABEL = {LOG_CORE, AIS_LOG_DOMAIN, "AIServiceImpl"};
    constexpr int MAX_MODEL_CREATE_NUM = 1; 
    constexpr double S_TO_MS = 1000.0; 
    const int YOLOV7_INPUT_HW = 640;

    inline static int Clamp(float val, int min, int max)
    {
        return val > min ? (val < max ? val : max) : min;
    }

    int HReverse(int h, LetterBox lb)
    {
        if (not lb.reverseAvailable)
        {
            int rh = Clamp(h, 0, lb.targetHeight);
            return rh;
        }
        int rh = Clamp(h, 0, lb.targetHeight) - lb.hPadTop;
        rh = Clamp(rh, 0, lb.targetHeight);
        rh = rh / lb.resizeScaleH;
        return rh;
    }

    int WReverse(int w, LetterBox lb)
    {
        if (not lb.reverseAvailable)
        {
            int rw = Clamp(w, 0, lb.targetWidth);
            return rw;
        }
        int rw = Clamp(w, 0, lb.targetWidth) - lb.wPadLeft;
        rw = Clamp(rw, 0, lb.targetWidth);
        rw = rw / lb.resizeScaleW;
        return rw;
    }
}


int32_t AIServiceImpl::Init(std::string& modelPath, int32_t* modelId)
{
    std::lock_guard<std::mutex> implLock(mutex_);
    int32_t ret = RETCODE_SUCCESS;
    if (modelPath.empty())
    {
        AIS_HILOGE("modelPath is invalid.");
        return RETCODE_INVALID_PARAM;
    }

    struct timeval t1 {};
    struct timeval t2 {};
    double costs = 0.0;
    gettimeofday(&t1, nullptr);

    if (!initFlag_)
    {
        ret = InitContext();
        if (ret != RETCODE_SUCCESS) {
            AIS_HILOGE("init failed, ret: %{public}d.", ret);
            return ret;
        }
        initFlag_ = true;
    }

    ret = CreateAndBuildModel(modelPath, modelId);
    if (ret != RETCODE_SUCCESS) {
        AIS_HILOGE("CreateAndBuildModel failed, ret: %{public}d.", ret);
        return ret;
    }
    gettimeofday(&t2, nullptr);
    costs = (double)(t2.tv_sec - t1.tv_sec) * S_TO_MS + (double)(t2.tv_usec - t1.tv_usec) / S_TO_MS;
    AIS_HILOGD("Init costs: %{public}f ms.", costs);

    AIS_HILOGD("Init success.");
    return RETCODE_SUCCESS;
}

int32_t AIServiceImpl::Process(const int32_t modelId, PicDesc& picDesc, std::string& picData, InferResult& result)
{
    std::lock_guard<std::mutex> implLock(mutex_);

    AIS_HILOGD("Process start.");
    if (modelId < 0 || modelMap_.find(modelId) == modelMap_.end())
    {
        AIS_HILOGE("modelId is invalid.");
        return RETCODE_INVALID_PARAM;
    }

    auto model = modelMap_[modelId];

    Mat blob;
    auto ret = PreProcess(picDesc, picData, blob);
    if (ret != RETCODE_SUCCESS)
    {
        return ret;
    }

    OH_AI_TensorHandleArray inputs = OH_AI_ModelGetInputs(model);
    if (inputs.handle_list == nullptr)
    {
        AIS_HILOGE("OH_AI_ModelGetInputs failed.");
        return RETCODE_MINDSPORE_MODEL_GET_INPUT_FAILED;
    }
    for (size_t i = 0; i < inputs.handle_num; ++i)
    {
        float *input_data = (float*)OH_AI_TensorGetMutableData(inputs.handle_list[i]);
        if (input_data == nullptr) {
            AIS_HILOGE("OH_AI_TensorGetMutableData failed.");
            return RETCODE_MINDSPORE_TENSOR_GET_MUTABLE_DATA_FAILED;
        }
        size_t tensorSize = OH_AI_TensorGetDataSize(inputs.handle_list[i]);
        memcpy(input_data, blob.data, tensorSize);
    }
    OH_AI_TensorHandleArray outputs = OH_AI_ModelGetOutputs(model);

    ret = Predict(model, inputs, outputs);
    if (inputs.handle_list == nullptr)
    {
        return ret;
    }

    PostProcess(model, outputs, result, picDesc);

    AIS_HILOGD("Process end.");
    return RETCODE_SUCCESS;
}

int32_t AIServiceImpl::DeInit()
{
    std::lock_guard<std::mutex> implLock(mutex_);
    AIS_HILOGI("DeInit start.");
    for(auto &it : modelMap_) {
        if (it.second != nullptr) {
            OH_AI_ModelDestroy(&it.second);
        }
    }
    modelMap_.clear();
    initFlag_ = false;
    modelIdTag_ = 0;
    return RETCODE_SUCCESS;
}

int32_t AIServiceImpl::InitContext()
{
    AIS_HILOGI("Init context start.");
    context_ = OH_AI_ContextCreate();
    if (context_ == nullptr)
    {
        AIS_HILOGE("OH_AI_ContextCreate failed.");
        return RETCODE_MINDSPORE_CREATE_CONTEXT_FAILED;
    }
    AIS_HILOGD("OH_AI_ContextCreate succeed.");

    // nnrt
    OH_AI_DeviceInfoHandle nnrt_device_info = OH_AI_CreateNNRTDeviceInfoByType(OH_AI_NNRTDEVICE_ACCELERATOR);
    if (nnrt_device_info == nullptr) {
        AIS_HILOGE("OH_AI_CreateNNRTDeviceInfoByType failed.");
        OH_AI_ContextDestroy(&context_);
        return OH_AI_STATUS_LITE_ERROR;
    }
    OH_AI_DeviceInfoSetPerformanceMode(nnrt_device_info, OH_AI_PERFORMANCE_HIGH);
    OH_AI_DeviceInfoSetEnableFP16(nnrt_device_info, false);
    OH_AI_ContextAddDeviceInfo(context_, nnrt_device_info);
    AIS_HILOGD("OH_AI_DeviceInfoCreate succeed.");

  return RETCODE_SUCCESS;   
}

int32_t AIServiceImpl::CreateAndBuildModel(std::string& modelPath, int32_t* modelId)
{
    if (modelIdTag_ > (MAX_MODEL_CREATE_NUM - 1)) {
        AIS_HILOGE("model create num is out of range.");
        return RETCODE_MINDSPORE_MODEL_NUM_OUT_OF_RANGE;
    }

    OH_AI_ModelHandle model = OH_AI_ModelCreate();
    if (model == nullptr)
    {
        AIS_HILOGE("OH_AI_ModelCreate failed.");
        return RETCODE_MINDSPORE_CREATE_MODEL_FAILED;
    }

    AIS_HILOGD("OH_AI_ModelCreate succeed.");

    AIS_HILOGD("Start to ModelBuildFromFile, model path: %{public}s", modelPath.c_str());
    auto ret = OH_AI_ModelBuildFromFile(model, modelPath.c_str(), OH_AI_MODELTYPE_MINDIR, context_);
    if (ret != OH_AI_STATUS_SUCCESS)
    {
        AIS_HILOGE("OH_AI_ModelBuildFromFile failed. ret: %{public}d", ret);
        OH_AI_ModelDestroy(&model);
        return RETCODE_MINDSPORE_BUILD_MODEL_FAILED;
    }
    modelMap_[modelIdTag_] = model;
    *modelId = modelIdTag_;
    modelIdTag_++;

    return RETCODE_SUCCESS; 
}

int32_t AIServiceImpl::PreProcess(PicDesc& picDesc, std::string& picData, cv::Mat& blob)
{
    AIS_HILOGD("PreProcess start.");

    if (picData.empty()) {
        AIS_HILOGE("PreProcess: input pic data is empty.");
        return RETCODE_INVALID_PARAM;
    }

    std::vector<char> buffer(picData.begin(), picData.end());
    Mat img = imdecode(buffer, 1);
    if (img.data == nullptr || img.channels() != 3 || img.rows == 0 || img.cols == 0) {
        AIS_HILOGE("PreProcess: imdecode failed.");
        return RETCODE_OPENCV_IMDECODE_FAILED;
    }

    img.convertTo(blob,CV_8UC3);
    cvtColor(blob, blob, cv::COLOR_BGR2RGB);
    
    return RETCODE_SUCCESS;
}

int32_t AIServiceImpl::Predict(OH_AI_ModelHandle& model, OH_AI_TensorHandleArray& inputs, OH_AI_TensorHandleArray& outputs)
{
    struct timeval t1 {};
    struct timeval t2 {};
    double costs = 0.0;
    gettimeofday(&t1, nullptr);
    auto ret = OH_AI_ModelPredict(model, inputs, &outputs, nullptr, nullptr);
    if (ret != OH_AI_STATUS_SUCCESS)
    {
        AIS_HILOGE("OH_AI_ModelPredict failed.");
        return RETCODE_MINDSPORE_MODEL_PREDICT_FAILED;
    }
    gettimeofday(&t2, nullptr);
    costs = (double)(t2.tv_sec - t1.tv_sec) * S_TO_MS + (double)(t2.tv_usec - t1.tv_usec) / S_TO_MS;
    AIS_HILOGD("OH_AI_ModelPredict costs: %{public}f ms.", costs);

    return RETCODE_SUCCESS;
}

void AIServiceImpl::PostProcess(OH_AI_ModelHandle& model, OH_AI_TensorHandleArray& outputs, InferResult& result, PicDesc& picDesc)
{
    std::vector<float> filterBoxes;
    std::vector<float> objProbs;
    std::vector<int> classId;
    int validCount = 0;
    for (size_t i = 0; i < outputs.handle_num; i++) {
        OH_AI_TensorHandle tensor = outputs.handle_list[i];
        int8_t *data =
            (int8_t *)OH_AI_TensorGetData(OH_AI_ModelGetOutputByTensorName(model, OH_AI_TensorGetName(tensor)));        
        int gridH = YOLOV7_OUTPUT_HW[i];
        int gridW = YOLOV7_OUTPUT_HW[i];
        int stride = YOLOV7_INPUT_HW / gridH;
        validCount = validCount + ProcessInt8(data, (int *)ANCHOR[i], gridH, gridW, stride, filterBoxes, 
                                                objProbs, classId, CONF_THRESHOLD, ZP[i], SCALE[i]);
    }

    result.count = 0;
    if (validCount <= 0) {
        AIS_HILOGW("yolo detect no object.");
        return;
    }

    std::vector<int> indexArray;
    for (int i = 0; i < validCount; ++i) {
        indexArray.push_back(i);
    }

    QuickSortIndiceInverse(objProbs, 0, validCount - 1, indexArray);

    std::set<int> class_set(std::begin(classId), std::end(classId));
    for (auto c : class_set) {
        NMS(validCount, filterBoxes, classId, indexArray, c, NMS_THRESHOLD);
    }

    int lastCount = 0;
    for (int i = 0; i < validCount; ++i) {
        if (indexArray[i] == -1 || objProbs[i] < CONF_THRESHOLD || lastCount >= OBJ_NUMB_MAX_SIZE) {
            continue;
        }
        int n = indexArray[i];
        result.objects[lastCount].left = filterBoxes[n * 4 + 0];
        result.objects[lastCount].top  = filterBoxes[n * 4 + 1];
        result.objects[lastCount].right  = result.objects[lastCount].left + filterBoxes[n * 4 + 2];
        result.objects[lastCount].bottom = result.objects[lastCount].top + filterBoxes[n * 4 + 3];
        result.objects[lastCount].name = CALSS_NAME[classId[n]];
        result.objects[lastCount].prop = objProbs[i];
        lastCount++;
    }
    result.count = lastCount;
    for (int i = 0; i < result.count; i++) {
        result.objects[i].left = WReverse(result.objects[i].left, picDesc.letterBox);
        result.objects[i].right = WReverse(result.objects[i].right, picDesc.letterBox);
        result.objects[i].top = HReverse(result.objects[i].top, picDesc.letterBox);
        result.objects[i].bottom = HReverse(result.objects[i].bottom, picDesc.letterBox);
    }

}
} // namespace AIS
} // namespace OHOS